import ast
import json
import javalang
import esprima

VERSION = "v2.3"
NODE_TYPE_DIRECTORY = "directory"
NODE_TYPE_FILE = "file"
NODE_TYPE_CLASS = "class"
NODE_TYPE_FUNCTION = "function"
EDGE_TYPE_CONTAINS = "contains"
EDGE_TYPE_INHERITS = "inherits"
EDGE_TYPE_INVOKES = "invokes"
EDGE_TYPE_IMPORTS = "imports"

VALID_NODE_TYPES = [
    NODE_TYPE_DIRECTORY,
    NODE_TYPE_FILE,
    NODE_TYPE_CLASS,
    NODE_TYPE_FUNCTION,
]
VALID_EDGE_TYPES = [
    EDGE_TYPE_CONTAINS,
    EDGE_TYPE_INHERITS,
    EDGE_TYPE_INVOKES,
    EDGE_TYPE_IMPORTS,
]


class PythonAnalyzer:
    def __init__(self, filename, source_code):
        self.filename = filename
        self.source_code = source_code
        self.source_lines = source_code.split("\n")
        self.nodes = []

    def analyze(self, tree):
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                self._visit_function(node)
            elif isinstance(node, ast.ClassDef):
                self._visit_class(node)
        return self.nodes

    def _extract_code_block(self, start_line, end_line):
        """提取指定行范围的代码"""
        if start_line <= 0 or start_line > len(self.source_lines):
            return ""

        if end_line is None or end_line <= 0:
            end_line = start_line

        # 调整为0基索引
        start_idx = start_line - 1
        end_idx = min(end_line, len(self.source_lines))

        return "\n".join(self.source_lines[start_idx:end_idx])

    def _visit_function(self, node):
        start_line = node.lineno
        end_line = node.end_lineno if hasattr(node, "end_lineno") else None
        code_block = self._extract_code_block(start_line, end_line)

        self.nodes.append(
            {
                "name": node.name,
                "parent_type": None,
                "type": NODE_TYPE_FUNCTION,
                "code": code_block,
                "start_line": start_line,
                "end_line": end_line,
            }
        )

    def _visit_class(self, node):
        start_line = node.lineno
        end_line = node.end_lineno if hasattr(node, "end_lineno") else None
        code_block = self._extract_code_block(start_line, end_line)

        self.nodes.append(
            {
                "name": node.name,
                "parent_type": None,
                "type": NODE_TYPE_CLASS,
                "code": code_block,
                "start_line": start_line,
                "end_line": end_line,
            }
        )

        # 访问类中的方法
        for child in node.body:
            if isinstance(child, ast.FunctionDef):
                method_start_line = child.lineno
                method_end_line = (
                    child.end_lineno if hasattr(child, "end_lineno") else None
                )
                method_code_block = self._extract_code_block(
                    method_start_line, method_end_line
                )

                self.nodes.append(
                    {
                        "name": f"{node.name}.{child.name}",
                        "parent_type": NODE_TYPE_CLASS,
                        "type": NODE_TYPE_FUNCTION,
                        "code": method_code_block,
                        "start_line": method_start_line,
                        "end_line": method_end_line,
                    }
                )


def analyze_file_py(filepath):
    with open(filepath, "r", encoding="utf-8") as file:
        code = file.read()

    try:
        if filepath.endswith(".py"):
            tree = ast.parse(code)
            analyzer = PythonAnalyzer(filepath, code)
            return analyzer.analyze(tree)
        elif filepath.endswith(".js"):
            tree = esprima.parseScript(code)
            return []  # JavaScript分析器待实现
        else:
            return []
    except Exception as e:
        print(f"Error parsing {filepath}: {e}")
        return []


if __name__ == "__main__":
    file_path = "/data_ext/ref_code/LocAgent/data/data/pp_data/fmxm/run.py"
    nodes = analyze_file(file_path)
    print(len(nodes))
    for node in nodes:
        print(json.dumps(node, ensure_ascii=False, indent=4))
